import requests
import json
import urllib3

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


# 获取所有ip白名单
def get_all_ip(url, header):
    url['uri'] = 'getsrciplist'
    url = "{http}://{domain}:{port}/{uri}".format(**url)
    print(url)
    result = requests.get(url=url, headers=header).json()
    return result


# 获取所有domain白名单
def get_all_domian(url, header):
    url['uri'] = 'getdstdomainlist'
    url = "{http}://{domain}:{port}/{uri}".format(**url)
    print(url)
    result = requests.get(url=url, headers=header).json()
    return result


# 添加ip白名单
def insert_ip(hosts, description, url, header):
    url['uri'] = 'insertsrcip'
    url = "{http}://{domain}:{port}/{uri}".format(**url)
    print(url)
    data = {
        "hosts": hosts,
        "description": description
    }
    print(json.dumps(data))
    requests.post(url=url, data=json.dumps(data), headers=header)


# 应用ip
def applyip(url, header):
    url['uri'] = 'applyip'
    url = "{http}://{domain}:{port}/{uri}".format(**url)
    result = requests.get(url=url, headers=header).json()
    return result


# 添加domain白名单
def insert_domain(dstdomain, url, header):
    url['uri'] = 'insertdomain'
    url = "{http}://{domain}:{port}/{uri}".format(**url)
    print(url)
    data = {
        "dstdomain": dstdomain
    }
    print(json.dumps(data))
    requests.post(url=url, data=json.dumps(data), headers=header)


# 应用domian
def applydomian(url, header):
    url['uri'] = 'applydomian'
    url = "{http}://{domain}:{port}/{uri}".format(**url)
    result = requests.get(url=url, headers=header).json()
    return result


def main():
    header = {
        'content-type': 'application/json'
    }
    url = {
        'http': 'http',
        'domain': '172.26.130.201',
        'port': 8000,

    }
    print("***squid代理服务器变更脚本***")
    print("---1、查看所有ip白名单")
    print("---2、查看所有domain白名单")
    print("---3、添加ip白名单")
    print("---4、添加domain白名单")

    choice = int(input("请输入选择："))
    if choice == 1:
        result = get_all_ip(url, header)
        for row in result:
            print(row)
    elif choice == 2:
        result = get_all_domian(url, header)
        for row in result:
            print(row)
    elif choice == 3:
        hosts = input("请输入ip(批量请用英文逗号隔开）:")
        description = input("请输入描述：")
        insert_ip(hosts, description, url, header)
        applyip(url, header)
        print("添加成功！")
    elif choice == 4:
        dstdomain = input("请输入domain(批量请用英文逗号隔开）:")
        insert_domain(dstdomain, url, header)
        applydomian(url, header)
        print("添加成功！")

    pass


if __name__ == '__main__':
    main()
